#!/usr/bin/env python3
"""
CardiacNet数据集诊断脚本

用于分析CardiacNet数据集的结构，找出空文件和问题文件
"""

import os
import sys
from pathlib import Path
from typing import Dict, List, Tuple
import json

def diagnose_cardiacnet(data_dir: str) -> Dict:
    """
    诊断CardiacNet数据集
    
    Args:
        data_dir: CardiacNet数据目录
        
    Returns:
        诊断结果
    """
    data_dir = Path(data_dir)
    
    if not data_dir.exists():
        print(f"错误：目录不存在: {data_dir}")
        return {}
    
    print(f"正在诊断CardiacNet数据集: {data_dir}")
    print("="*60)
    
    # 查找所有.nii和.nii.gz文件
    nii_files = list(data_dir.rglob('*.nii'))
    nii_gz_files = list(data_dir.rglob('*.nii.gz'))
    
    print(f"\n找到文件:")
    print(f"  .nii文件: {len(nii_files)} 个")
    print(f"  .nii.gz文件: {len(nii_gz_files)} 个")
    
    # 分析.nii文件
    empty_nii = []
    small_nii = []
    valid_nii = []
    dir_nii = []  # 被误当作目录的.nii路径
    files_in_dirs = []  # 目录内的实际文件
    
    for f in nii_files:
        if f.is_file():
            try:
                size = f.stat().st_size
                if size == 0:
                    empty_nii.append((str(f), size))
                elif size < 1024:
                    small_nii.append((str(f), size))
                else:
                    valid_nii.append((str(f), size))
            except Exception:
                pass
        elif f.is_dir():
            dir_nii.append(str(f))
            # 检查目录内是否有实际的文件
            inner_files = list(f.rglob('*'))
            inner_nii = [x for x in inner_files if x.is_file() and x.suffix.lower() in ['.nii', '.gz']]
            if inner_nii:
                files_in_dirs.extend(inner_nii)
                print(f"  发现目录 {f.name} 内包含 {len(inner_nii)} 个文件")
    
    # 分析.nii.gz文件
    empty_gz = []
    small_gz = []
    valid_gz = []
    
    for f in nii_gz_files:
        if f.is_file():
            try:
                size = f.stat().st_size
                if size == 0:
                    empty_gz.append((str(f), size))
                elif size < 1024:
                    small_gz.append((str(f), size))
                else:
                    valid_gz.append((str(f), size))
            except Exception:
                pass
    
    # 分析目录内的文件
    valid_in_dirs = []
    for f in files_in_dirs:
        try:
            if f.is_file():
                size = f.stat().st_size
                if size > 1024:
                    valid_in_dirs.append((str(f), size))
        except Exception:
            pass
    
    # 打印统计
    print(f"\n.nii文件分析:")
    print(f"  有效文件: {len(valid_nii)} 个")
    print(f"  空文件: {len(empty_nii)} 个")
    print(f"  过小文件: {len(small_nii)} 个")
    print(f"  目录（非文件）: {len(dir_nii)} 个")
    if files_in_dirs:
        print(f"  目录内的文件: {len(files_in_dirs)} 个")
        print(f"  目录内的有效文件: {len(valid_in_dirs)} 个")
    
    print(f"\n.nii.gz文件分析:")
    print(f"  有效文件: {len(valid_gz)} 个")
    print(f"  空文件: {len(empty_gz)} 个")
    print(f"  过小文件: {len(small_gz)} 个")
    
    # 检查是否有对应的.nii.gz文件可以替代空的.nii文件
    if empty_nii or small_nii:
        print(f"\n尝试查找对应的.nii.gz文件...")
        replacements = []
        for nii_path, _ in empty_nii + small_nii:
            nii_file = Path(nii_path)
            # 检查是否有对应的.nii.gz文件
            gz_path = nii_file.parent / f"{nii_file.name}.gz"
            if gz_path.exists() and gz_path.is_file() and gz_path.stat().st_size > 1024:
                replacements.append((nii_path, str(gz_path)))
        
        if replacements:
            print(f"  找到 {len(replacements)} 个可用的.nii.gz替代文件")
            print(f"\n建议：使用.nii.gz文件替代空的.nii文件")
            for nii, gz in replacements[:5]:  # 只显示前5个
                print(f"    {Path(nii).name} -> {Path(gz).name}")
            if len(replacements) > 5:
                print(f"    ... 还有 {len(replacements) - 5} 个")
    
    # 生成修复建议
    suggestions = []
    
    if len(valid_gz) > len(valid_nii):
        suggestions.append(f"建议：优先使用.nii.gz文件（找到{len(valid_gz)}个有效文件）")
    
    if dir_nii:
        if valid_in_dirs:
            suggestions.append(f"发现{len(dir_nii)}个.nii目录，其中包含{len(valid_in_dirs)}个有效文件，将自动处理")
        else:
            suggestions.append(f"发现{len(dir_nii)}个.nii目录，可能需要检查数据集结构")
    
    if empty_nii:
        suggestions.append(f"发现{len(empty_nii)}个空.nii文件，这些文件将被自动跳过")
    
    if suggestions:
        print(f"\n修复建议:")
        for i, suggestion in enumerate(suggestions, 1):
            print(f"  {i}. {suggestion}")
    
    # 保存诊断结果
    diagnosis_result = {
        'data_dir': str(data_dir),
        'nii_files': {
            'total': len(nii_files),
            'valid': len(valid_nii),
            'empty': len(empty_nii),
            'small': len(small_nii),
            'directories': len(dir_nii)
        },
        'nii_gz_files': {
            'total': len(nii_gz_files),
            'valid': len(valid_gz),
            'empty': len(empty_gz),
            'small': len(small_gz)
        },
        'replacements': len(replacements) if 'replacements' in locals() else 0,
        'suggestions': suggestions
    }
    
    # 保存到文件
    output_file = data_dir / "cardiacnet_diagnosis.json"
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(diagnosis_result, f, indent=2, ensure_ascii=False)
    
    print(f"\n诊断结果已保存到: {output_file}")
    
    return diagnosis_result


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="CardiacNet数据集诊断工具")
    parser.add_argument("--cardiacnet_dir", type=str, default="D:/Data/CardiacNet",
                       help="CardiacNet数据目录")
    
    args = parser.parse_args()
    
    result = diagnose_cardiacnet(args.cardiacnet_dir)
    
    if result:
        print("\n" + "="*60)
        print("诊断完成！")
        print("="*60)


if __name__ == "__main__":
    main()

